"""
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    created by lei.chang on '20/11/2024'
    comment: agent智能相关
"""
import uuid
import ujson
from sqlalchemy import func, distinct
from datetime import datetime

from config import Config
from project.lib.decorator import with_async_redis
from project.core.base import BodyParams, Select
from project.core.function import serializer
from project.core.exception import BizError, NoRecordError
from project.model.basic import User, Device, FlushDeviceMode
from project.model.database.cariabot import Organization
from project.model.database.node import Node
from project.model.database.knowledge import Knowledge, KnowledgeDetail, KnowledgeData
from project.model.database.agent import AgentKnowledgeRelation, Chat, ChatMessage
from project.model.basic import Operator
from project.model.database.agent import Agent
from project.micro_service.apollo import Apollo
from project.micro_service.zeus import Zeus
from project.service.flow import FlowService


class AgentService:

    @staticmethod
    async def get_global_intent(session, device: Device, user_message, env='normal', memories_record=True):
        """
        :param session:
        :param device: 设备
        :param user_message:
        :param env: 指定模型
        :param memories_record: 是否添加记忆
        :return:
        """
        user = User(dev_id=device.dev_id, bot_id='udi', user_id='uu01')

        # 搜集建议
        if user_message:
            await Apollo.collect_advice(user, user_message)

        # 取得agent_profile, user_role
        organization = await Organization.fetch_one(session, wheres=(Organization.code == device.organization_code))
        agent_profile = organization.agent_profile if organization else '你是一个在酒店房间里的人工智能机器人'
        user_role = organization.user_role if organization else '客人'

        flows = await FlowService.get_cariabot_flow_list(session, device)
        for flow in flows:
            flow.title = flow.name
            flow.description = flow.introduction
        flow_map = {v.name: v for v in flows}

        # 根据设备编号取得已绑定的知识库
        knowledge_detail_codes = None
        if device.agent_code:
            agent = await Agent.fetch_one(session, wheres=(Agent.code == device.agent_code, ~Agent.is_deleted))
            if agent:
                knowledge_details = await Select(selects=(KnowledgeDetail),
                                                 joins=((KnowledgeData, KnowledgeData.id == KnowledgeDetail.knowledge_data_id), ),
                                                 wheres=(KnowledgeData.knowledge_id.in_(agent.knowledge_ids),
                                                         ~KnowledgeData.is_deleted,
                                                         ~KnowledgeDetail.is_deleted)).fetch_all(session)
                knowledge_detail_codes = [v.code for v in knowledge_details]

        params = dict(
            user=user,
            user_message=user_message,
            flows=[v.dict('code', 'title', 'description', 'matching_intent') for v in flows],
            agent_profile=agent_profile,
            user_role=user_role,
            memories_record=memories_record,
            organization_code=device.organization_code,
            knowledge_detail_codes=knowledge_detail_codes
        )
        if env == 'demo':
            params['chat_model_name'] = 'deepseek-chat'

        # 取得全局意图
        intent = await Apollo.get_global_intent(**params)

        if intent.get('intent') == 'service':
            flow = flow_map.get(intent.get('action'))

            if flow:
                intent['flow_code'] = flow.code if flow else ''
                node = await Node.fetch_one(session, wheres=(Node.id == flow.id, Node.is_first, ~Node.is_deleted))
                intent['node_code'] = node.code if node else ''

        elif intent.get('intent') == 'report' and organization and organization.sleep_agent_message:
            # 添加记忆 - 睡眠指定话术
            await Apollo.add_user_memories(user, agent_profile, user_role, agent_message=organization.sleep_agent_message)

        return intent

    @staticmethod
    @with_async_redis()
    async def add_agent(session, operator: Operator, body_params: BodyParams, rds=None):
        """ 添加智能体
        :param session:
        :param operator:
        :param body_params:
        :param rds:
        :return:
        """
        agent = await Agent.create_modify(session, operator_code=operator.code, organization_code=operator.organization_code, **body_params.dict())
        if not agent:
            raise BizError('智能体添加失败')

        default_agent = await Agent.fetch_one(session, wheres=(Agent.organization_code == operator.organization_code, Agent.is_default, ~Agent.is_deleted))

        now = datetime.now()
        index = await rds.incr(f"{Config.PROJECT_NAME}:aget_code_index")
        agent.code = f'{now.strftime("%Y")[-2:]}{str(now.isocalendar()[1]).zfill(2)}{str(now.weekday())}{index}'
        if not default_agent:
            agent.is_default = True
            await Zeus.switch_device_agent(operator.organization_code, '', agent.code)

        await session.commit()

        return agent

    @staticmethod
    async def update_agent(session, body_params: BodyParams):
        """ 修改智能体
        :param session:
        :param body_params:
        :return:
        """
        agent = await Agent.fetch_one(session, wheres=(Agent.id == body_params.id))
        if not agent:
            raise BizError('智能体不存在')

        agent.set_model_value(**body_params.dict())
        agent.bind_knowledge_count = len(agent.knowledge_ids)

        bind_knowledge_ids = [v for v in agent.knowledge_ids]
        # 清空智能体和知识库的绑定关系
        data = await AgentKnowledgeRelation.fetch_all(session, wheres=(AgentKnowledgeRelation.agent_id == agent.id))
        for v in data:
            bind_knowledge_ids.append(v.knowledge_id)
            await session.delete(v)

        # 添加绑定关系
        for knowledge_id in agent.knowledge_ids:
            await AgentKnowledgeRelation.create_modify(session, agent_id=agent.id, knowledge_id=knowledge_id, commit=False)

        # 更新知识库的关联agent数量
        data = await Select(selects=(AgentKnowledgeRelation.knowledge_id, func.count(distinct(AgentKnowledgeRelation.agent_id))),
                            wheres=(AgentKnowledgeRelation.knowledge_id.in_(bind_knowledge_ids)),
                            groups=(AgentKnowledgeRelation.knowledge_id,)).fetch_all(session)
        bind_agent_count_map = {knowledge_id: bind_agent_count for knowledge_id, bind_agent_count in data}
        knowledges = await Knowledge.fetch_all(session, wheres=(Knowledge.id.in_(bind_knowledge_ids)))
        for knowledge in knowledges:
            knowledge.bind_agent_count = bind_agent_count_map.get(knowledge.id, 0)

        await session.commit()

    @staticmethod
    async def delete_agent(session, agent_id):
        """ 删除智能体
        :param session:
        :param agent_id:
        :return:
        """
        agent = await Agent.fetch_one(session, wheres=(Agent.id == agent_id))
        if not agent:
            raise BizError('智能体不存在')
        if agent.is_default:
            raise BizError('默认智能体无法删除')
        devices = await Zeus.get_device_list(agent_code=agent.code, is_internal=1, count=1)
        if devices:
            raise BizError('智能体已被设备绑定，无法删除')
        agent.is_deleted = True

        bind_knowledge_ids = [v for v in agent.knowledge_ids]
        # 清空智能体和知识库的绑定关系
        data = await AgentKnowledgeRelation.fetch_all(session, wheres=(AgentKnowledgeRelation.agent_id == agent.id))
        for v in data:
            bind_knowledge_ids.append(v.knowledge_id)
            await session.delete(v)

        # 更新知识库的关联agent数量
        data = await Select(selects=(AgentKnowledgeRelation.knowledge_id, func.count(distinct(AgentKnowledgeRelation.agent_id))),
                            wheres=(AgentKnowledgeRelation.knowledge_id.in_(bind_knowledge_ids)),
                            groups=(AgentKnowledgeRelation.knowledge_id,)).fetch_all(session)
        bind_agent_count_map = {knowledge_id: bind_agent_count for knowledge_id, bind_agent_count in data}
        knowledges = await Knowledge.fetch_all(session, wheres=(Knowledge.id.in_(bind_knowledge_ids)))
        for knowledge in knowledges:
            knowledge.bind_agent_count = bind_agent_count_map.get(knowledge.id, 0)

        await session.commit()

    @staticmethod
    def get_conversation_key(conversation_id):
        """
        :param conversation_id:
        :return:
        """
        return f'{Config.PROJECT_NAME}:agent_conversation_{conversation_id}'

    @staticmethod
    @with_async_redis()
    async def create_conversation(session, agent_id, rds=None):
        """ 创建智能体对话
        :param session:
        :param agent_id:
        :param rds:
        :return:
        """
        agent = await Agent.fetch_one(session, wheres=(Agent.id == agent_id, ~Agent.is_deleted))
        if not agent:
            raise BizError('智能体不存在')

        conversation_id = str(uuid.uuid4()).replace('-', '')
        key = AgentService.get_conversation_key(conversation_id)
        await rds.hset(key, 'agent_id', agent_id)
        await rds.hset(key, 'messages', ujson.dumps([]))

        return conversation_id

    @staticmethod
    @with_async_redis()
    async def update_conversation(session, conversation_id, body_params, rds=None):
        """ 修改智能体对话
        :param session:
        :param conversation_id: 新会话
        :param body_params:
        :param rds:
        :return:
        """
        key = AgentService.get_conversation_key(conversation_id)
        # 新会话
        if await rds.exists(key):
            for k, v in body_params.dict().items():
                await rds.hset(key, k, v)
        else:
            try:
                await Chat.create_modify(session, wheres=(Chat.conversation_id == conversation_id), no_record_error=True, **body_params.dict())
            except NoRecordError:
                raise BizError('会话不存在')

    @staticmethod
    @with_async_redis()
    async def do_agent_chat_streaming(session, conversation_id, user_message, rds=None):
        """ 智能体对话(流式)
        :param session:
        :param conversation_id:
        :param user_message:
        :param rds:
        :return:
        """
        key = AgentService.get_conversation_key(conversation_id)
        # 新会话
        if await rds.exists(key):
            agent_id = await rds.hget(key, 'agent_id')
            agent = await Agent.fetch_one(session, wheres=(Agent.id == agent_id, ~Agent.is_deleted))
            if not agent:
                yield '智能体不存在'
                return

            # 取得关联知识库编号列表
            knowledge_details = await KnowledgeDetail.fetch_all(session, wheres=(KnowledgeDetail.knowledge_id.in_(agent.knowledge_ids), ~KnowledgeDetail.is_deleted))
            knowledge_detail_codes = [v.code for v in knowledge_details]

            # 取得历史对话信息
            messages = await rds.hget(key, 'messages')
            messages = ujson.loads(messages if messages else '[]')

            chat_start_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            response = ''
            async for content in await Apollo.do_agent_chat_streaming(knowledge_detail_codes, user_message, chat_messages=messages):
                response += content
                yield ujson.dumps({'response': content}, ensure_ascii=False)

            # 添加记忆
            messages.append({'role': 'user', 'message': user_message, 'create_time': chat_start_time})
            messages.append({'role': 'agent', 'message': response, 'create_time': datetime.now().strftime('%Y-%m-%d %H:%M:%S')})
            await rds.hset(key, 'messages', ujson.dumps(messages, ensure_ascii=False))

        # 历史会话
        else:
            chat = await Chat.fetch_one(session, wheres=(Chat.conversation_id == conversation_id, ~Chat.is_deleted))
            if not chat:
                yield '会话不存在'
                return
            agent = await Agent.fetch_one(session, wheres=(Agent.id == chat.agent_id, ~Agent.is_deleted))
            if not agent:
                yield '智能体不存在'
                return

            # 取得关联知识库编号列表
            knowledge_details = await KnowledgeDetail.fetch_all(session, wheres=(KnowledgeDetail.knowledge_id.in_(agent.knowledge_ids), ~KnowledgeDetail.is_deleted))
            knowledge_detail_codes = [v.code for v in knowledge_details]

            # 取得历史对话信息
            data = await ChatMessage.fetch_all(session, wheres=(ChatMessage.chat_id == chat.id, ~ChatMessage.is_deleted), orders=(ChatMessage.id,))
            messages = [serializer(v.dict('role', 'message', 'create_time')) for v in data]

            await ChatMessage.create_modify(session, role='role', message=user_message)
            response = ''
            async for content in await Apollo.do_agent_chat_streaming(knowledge_detail_codes, user_message, chat_messages=messages):
                response += content
                yield ujson.dumps({'response': content}, ensure_ascii=False)

            await ChatMessage.create_modify(session, role='agent', message=response)

    @staticmethod
    @with_async_redis()
    async def save_conversation(session, conversation_id, rds=None):
        """ 保存会话
        :param session:
        :param conversation_id:
        :param rds:
        :return:
        """
        chat = await Chat.fetch_one(session, wheres=(Chat.conversation_id == conversation_id, ~Chat.is_deleted))
        if chat:
            raise BizError('会话已保存')

        key = AgentService.get_conversation_key(conversation_id)
        params = await rds.hgetall(key)
        if not params.get('title'):
            index = await rds.incr(f"{Config.PROJECT_NAME}:agent_conversation_index")
            params['title'] = f'未命名对话({str(index)})'
        params.pop('conversation_id', None)
        chat = await Chat.create_modify(session, conversation_id=conversation_id, **params, commit=False)
        if not chat:
            raise BizError('会话保存失败')

        messages = await rds.hget(key, 'messages')
        messages = ujson.loads(messages)
        for v in messages:
            await ChatMessage.create_modify(session, chat_id=chat.id, role=v['role'], message=v['message'], commit=False)
        chat.message_count = len(messages)

        # 删除会话缓存
        await rds.delete(key)

        await session.commit()

    @staticmethod
    @with_async_redis()
    async def delete_conversation(session, conversation_id, rds=None):
        """ 删除会话
        :param session:
        :param conversation_id:
        :param rds:
        :return:
        """
        key = AgentService.get_conversation_key(conversation_id)
        # 删除缓存会话
        if await rds.exists(key):
            await rds.delete(key)

        chat = await Chat.fetch_one(session, wheres=(Chat.conversation_id == conversation_id, ~Chat.is_deleted))
        if chat:
            chat.is_deleted = True
            messages = await ChatMessage.fetch_all(session, wheres=(ChatMessage.chat_id == chat.id, ~ChatMessage.is_deleted))
            for message in messages:
                message.is_deleted = True

        await session.commit()

    @staticmethod
    async def set_default_agent(session, agent_id, flush_device_mode):
        """ 设置默认机构默认智能体
        :param session:
        :param agent_id:
        :param flush_device_mode:
        :return:
        """
        agent = await Agent.fetch_one(session, wheres=(Agent.id == agent_id, ~Agent.is_deleted))
        if not agent:
            raise BizError('智能体不存在')

        # 原默认智能体
        origin_default_agent = await Agent.fetch_one(session, wheres=(Agent.organization_code == agent.organization_code, ~Agent.is_deleted, Agent.is_default))

        # 更新默认智能体设备
        if flush_device_mode.value == str(FlushDeviceMode.default_device_only.value):
            await Zeus.switch_device_agent(agent.organization_code, origin_default_agent.code if origin_default_agent else '', agent.code)
        elif flush_device_mode.value == str(FlushDeviceMode.whole_organization.value):
            await Zeus.switch_device_agent(agent.organization_code, '', agent.code)

        agent.is_default = True
        if origin_default_agent:
            origin_default_agent.is_default = False

        await session.commit()

